"""
Demo 02: 统计特征向量（均值、方差、偏度、峰度等）
从时间序列窗口中提取统计特征作为向量表示
"""

import numpy as np
import scipy.stats as st
from demo_01_raw import RawVectorizer


class StatisticalVectorizer:
    """统计特征向量化处理类"""
    
    def __init__(self, window_size=60, step_size=5):
        """
        初始化参数
        
        Args:
            window_size: 窗口大小
            step_size: 步长
        """
        self.window_size = window_size
        self.step_size = step_size
        self.raw_vectorizer = RawVectorizer(window_size, step_size)
        
    def extract_statistical_features(self, windows):
        """
        从窗口中提取统计特征
        
        Args:
            windows: 窗口矩阵 (n_windows, window_size)
            
        Returns:
            numpy.ndarray: 统计特征矩阵 (n_windows, n_features)
        """
        # 计算各种统计特征
        mean = windows.mean(axis=1, keepdims=True)
        std = windows.std(axis=1, keepdims=True)
        max_val = windows.max(axis=1, keepdims=True)
        min_val = windows.min(axis=1, keepdims=True)
        skew = st.skew(windows, axis=1, keepdims=True)
        kurt = st.kurtosis(windows, axis=1, keepdims=True)
        
        # 额外的统计特征
        median = np.median(windows, axis=1, keepdims=True)
        q25 = np.percentile(windows, 25, axis=1, keepdims=True)
        q75 = np.percentile(windows, 75, axis=1, keepdims=True)
        iqr = q75 - q25  # 四分位距
        
        # 计算一阶差分的统计量
        diff = np.diff(windows, axis=1)
        diff_mean = diff.mean(axis=1, keepdims=True)
        diff_std = diff.std(axis=1, keepdims=True)
        
        # 合并所有特征
        features = np.hstack([
            mean, std, max_val, min_val, skew, kurt,
            median, q25, q75, iqr, diff_mean, diff_std
        ])
        
        return features
    
    def get_feature_names(self):
        """
        获取特征名称列表
        
        Returns:
            list: 特征名称列表
        """
        return [
            'mean', 'std', 'max', 'min', 'skew', 'kurtosis',
            'median', 'q25', 'q75', 'iqr', 'diff_mean', 'diff_std'
        ]
    
    def fit_transform(self, ticker="AAPL", start="2020-01-01", end="2024-12-31"):
        """
        完整的统计特征向量化流程
        
        Args:
            ticker: 股票代码
            start: 开始日期
            end: 结束日期
            
        Returns:
            tuple: (原始窗口矩阵, 统计特征矩阵)
        """
        # 获取原始窗口
        X_raw = self.raw_vectorizer.fit_transform(ticker, start, end)
        
        # 提取统计特征
        X_stats = self.extract_statistical_features(X_raw)
        
        return X_raw, X_stats
    
    def transform(self, windows):
        """
        对给定窗口提取统计特征
        
        Args:
            windows: 窗口矩阵
            
        Returns:
            numpy.ndarray: 统计特征矩阵
        """
        return self.extract_statistical_features(windows)
    
    def describe_features(self, features):
        """
        描述统计特征
        
        Args:
            features: 特征矩阵
        """
        feature_names = self.get_feature_names()
        print("\n特征统计摘要:")
        print("-" * 60)
        print(f"{'Feature':<15} {'Mean':>10} {'Std':>10} {'Min':>10} {'Max':>10}")
        print("-" * 60)
        
        for i, name in enumerate(feature_names):
            feat_col = features[:, i]
            print(f"{name:<15} {feat_col.mean():>10.4f} {feat_col.std():>10.4f} "
                  f"{feat_col.min():>10.4f} {feat_col.max():>10.4f}")


def main():
    """主函数，演示使用方法"""
    print("=" * 60)
    print("Demo 02: 统计特征向量")
    print("=" * 60)
    
    # 创建统计特征向量化器
    vectorizer = StatisticalVectorizer(window_size=60, step_size=5)
    
    # 执行向量化
    print("正在下载AAPL股票数据并提取统计特征...")
    X_raw, X_stats = vectorizer.fit_transform(
        ticker="AAPL",
        start="2020-01-01",
        end="2024-12-31"
    )
    
    # 输出结果
    print(f"\n原始窗口 shape: {X_raw.shape}")
    print(f"统计特征向量 shape: {X_stats.shape}")
    print(f"特征数量: {X_stats.shape[1]}")
    
    # 显示特征名称
    feature_names = vectorizer.get_feature_names()
    print(f"\n特征列表: {feature_names}")
    
    # 显示第一个样本的特征值
    print(f"\n第一个窗口的统计特征:")
    for i, name in enumerate(feature_names):
        print(f"  {name}: {X_stats[0, i]:.4f}")
    
    # 特征统计摘要
    vectorizer.describe_features(X_stats)
    
    return X_raw, X_stats


if __name__ == "__main__":
    X_raw, X_stats = main()
